//#include <stdio.h>
//#include <stdint.h>
//#include <stdlib.h>
//#include <sys/types.h>
//#include <limits.h>
#include <float.h>
//#include <libgen.h>
//#include <ctype.h>
#include <getopt.h>
//#include <math.h>
//#include <assert.h>

#include <cuqu/cuqu.h>
#ifdef CUPRINTF_INCLUDED
#include <utils/cuPrintf.cu>
#endif
#include <cuqu/detail/barrier_oneloop.h>
#include <cuqu/detail/barrier_lockfree.h>
#include <cuqu/detail/barrier_tree_based.h>


// Defines default
#define DEFAULT_NUM_BLOCKS      6
#define DEFAULT_NUM_THREADS     512
#define DEFAULT_NUM_WORKERS     (DEFAULT_NUM_BLOCKS * DEFAULT_NUM_THREADS)    
#define DEFAULT_TEST_REPEATS    100
#define DEFAULT_KERN_CALLS      100
#define DEFAULT_BLOCK_STEP      1
#define DEFAULT_THREAD_STEP     64
#define DEFAULT_WORKER_STEP     128
#define DEFAULT_SYNC_METHOD     OneLoop
//#define CUPRINTF_CU


// Synchronization implementations
typedef enum { 
    NoSync,         // No synchronization
    CPUImplicit,    // CPU implicit synchronization
    CPUExplicit,    // CPU explicit synchronization
    OneLoop,        // GPU one loop synchronization
    LockFree,       // GPU lock free synchronization
    TreeBased       // GPU tree based synchronization    
} SyncMethod;

// User parameters
typedef struct {
    bool autoBlocks;    // set auto test blocks
    bool autoThreads;   // set auto test threads
    bool autoWorkers;   // set auto test working threads (workers)
    int  blockStep;     // step block
    int  threadStep;    // step thread
    int  workerStep;    // step workers
    int  numBlocks;     // number of blocks
    int  numThreads;    // number of threads
    int  numWorkers;    // number of workers
    int  kernelCalls;   // number of kernel loops
    int  testRepeats;   // number of test loops
    SyncMethod syncMethod;    // barrier method to test
} UserParameters;

// Test measures
typedef struct {
    float averageTime;  // test average time 
    float minTime;      // minimum time detected
    float maxTime;      // maximum time detected
    float errorTime;    // standard deviation 
} Measures;

// Program name used in printUsage function
char *programName = NULL;

// Host functions
void printUsage(bool);
void parseArguments(int, char**, UserParameters*);
void parseNumberGTZ(char, char*, int*);
void getMeasures(float[], size_t, Measures*);

////////////////////////////////////////////////////////////////////////////////
//
// Device code
//
////////////////////////////////////////////////////////////////////////////////

// Kernel using GPU barrier synchronization
template <typename BarrierType> 
__global__ void kernelGPU(float *arrayOut, float *arrayIn1, float *arrayIn2, int numWorkers, int numCalls, BarrierType *barrier) {    

    const unsigned int threadID  = threadIdx.x + blockDim.x * threadIdx.y + blockDim.x * blockDim.y * threadIdx.z;
    const unsigned int blockID   = blockIdx.x + gridDim.x * blockIdx.y;
    const unsigned int blockSize = blockDim.x * blockDim.y * blockDim.z;
    const unsigned int gridSize  = gridDim.y * gridDim.x;
    const unsigned int globalID  = threadID + blockID * blockSize;
    
    for (int n = 0; n < numCalls; n++) {  
        if (globalID < numWorkers) {                                     
            arrayOut[globalID] = arrayIn1[globalID] * n + arrayIn2[globalID] * (numCalls - n);    
        }                                                           
        barrier_wait(barrier, gridSize);          
    }                                                               
}

// Kernel without sync
__global__ void kernelNoSync(float *arrayOut, float *arrayIn1, float *arrayIn2, int numWorkers, int numCalls) {

    const unsigned int threadID  = threadIdx.x + blockDim.x * threadIdx.y + blockDim.x * blockDim.y * threadIdx.z;
    const unsigned int blockID   = blockIdx.x + gridDim.x * blockIdx.y;
    const unsigned int blockSize = blockDim.x * blockDim.y * blockDim.z;
    //const unsigned int gridSize  = gridDim.y * gridDim.x;
    const unsigned int globalID  = threadID + blockID * blockSize;

    for (int n = 0; n < numCalls; n++) {
        if (globalID < numWorkers) {
            arrayOut[globalID] = arrayIn1[globalID] * n + arrayIn2[globalID] * (numCalls - n);
        }
    }
}


// Kernel using CPU implicit synchronization
__global__ void kernelCPU(float *arrayOut, float *arrayIn1, float *arrayIn2, int numWorkers) {    

    const unsigned int threadID  = threadIdx.x + blockDim.x * threadIdx.y + blockDim.x * blockDim.y * threadIdx.z;
    const unsigned int blockID   = blockIdx.x + gridDim.x * blockIdx.y;
    const unsigned int blockSize = blockDim.x * blockDim.y * blockDim.z;
    //const unsigned int gridSize  = gridDim.y * gridDim.x;
    const unsigned int globalID  = threadID + blockID * blockSize;
        
    if (globalID < numWorkers) {
        arrayOut[globalID] = arrayIn1[globalID] * numWorkers + arrayIn2[globalID] * (numWorkers - blockSize);
    }                                                           
}

// Device barriers
cuqu::detail::barrier_oneloop *barrierOneLoop; 
cuqu::detail::barrier_lockfree *barrierLockFree;
cuqu::detail::barrier_tree_based *barrierTreeBased;


////////////////////////////////////////////////////////////////////////////////
//
// Switch kernel
//
////////////////////////////////////////////////////////////////////////////////
void SwitchKernel(float *arrayOut, float *arrayIn1, float *arrayIn2, int numBlocks, int numThreads, int numWorkers, int numCalls, SyncMethod syncMethod) {

    switch (syncMethod) {

        case NoSync:           
            kernelNoSync<<<numBlocks, numThreads>>>(arrayOut, arrayIn1, arrayIn2, numWorkers, numCalls);	
            break;

        case CPUImplicit:
            for (int n = 0; n < numCalls; n++) {
                kernelCPU<<<numBlocks, numThreads>>>(arrayOut, arrayIn1, arrayIn2, numWorkers);
            } 
            break;
        
        case CPUExplicit:
            for (int n = 0; n < numCalls; n++) {
                kernelCPU<<<numBlocks, numThreads>>>(arrayOut, arrayIn1, arrayIn2, numWorkers);
                cudaThreadSynchronize();
            } 
            break;

        case OneLoop:
            kernelGPU<<<numBlocks,numThreads>>>(arrayOut, arrayIn1, arrayIn2, numWorkers, numCalls, barrierOneLoop);
            break;

        case LockFree:
            kernelGPU<<<numBlocks,numThreads>>>(arrayOut, arrayIn1, arrayIn2, numWorkers, numCalls, barrierLockFree);
            break;

        case TreeBased:
            kernelGPU<<<numBlocks,numThreads>>>(arrayOut, arrayIn1, arrayIn2, numWorkers, numCalls, barrierTreeBased);
            break;
    }
}


// Estimate barrier time
void barrierTime(int numBlocks, int numThreads, int numWorkers, int numCalls, int numTests, SyncMethod syncMethod, Measures* measures) {

    // Start and stop CUDA events
    cudaEvent_t start, stop;
    float times[numTests];
    int numElements = numBlocks * numThreads;
    float *d_arrayOut, *d_arrayIn1, *d_arrayIn2;

    // Allocate device memory
    cuquSafeCall(cudaMalloc(&d_arrayOut, sizeof(float) * numElements));
    cuquSafeCall(cudaMalloc(&d_arrayIn1, sizeof(float) * numElements));
    cuquSafeCall(cudaMalloc(&d_arrayIn2, sizeof(float) * numElements));

    // Initialize device memory
    cuquSafeCall(cudaMemset(d_arrayOut, 0, sizeof(float) * numElements));
    cuquSafeCall(cudaMemset(d_arrayIn1, 0, sizeof(float) * numElements));
    cuquSafeCall(cudaMemset(d_arrayIn2, 0, sizeof(float) * numElements));

    // Create start and stop events
    cuquSafeCall(cudaEventCreate(&start));
    cuquSafeCall(cudaEventCreate(&stop));

    for (int i = 0; i < numTests; i++) { 
        
        // Record start event
        cuquSafeCall(cudaEventRecord(start, 0));

        // Call kernel
        SwitchKernel(d_arrayOut, d_arrayIn1, d_arrayIn2, numBlocks, numThreads, numWorkers, numCalls, syncMethod);
     
        // Record stop event and synchronize events
        cuquSafeCall(cudaEventRecord(stop, 0));
        cuquSafeCall(cudaEventSynchronize(stop));

        // Calculate elapsed time
        cuquSafeCall(cudaEventElapsedTime(&(times[i]), start, stop));

        // Estimate synchronization time
        times[i] = times[i] / (float) numCalls;
    }

    // Destroy start and stop events
    cuquSafeCall(cudaEventDestroy(start));
    cuquSafeCall(cudaEventDestroy(stop));

    // Free device memory
    cuquSafeCall(cudaFree(d_arrayOut));
    cuquSafeCall(cudaFree(d_arrayIn1));
    cuquSafeCall(cudaFree(d_arrayIn2));

    // Get measures
    getMeasures(times, numTests, measures);
}


// Run the test
void runTest(UserParameters *params) {

    Measures measures;
    int startBlocks = params->numBlocks;
    int startThreads = params->numThreads;
    int maxBlocks = params->numBlocks;
    int maxThreads = params->numThreads;

    // Init device resource
    cuqu::detail::barrier_oneloop h_barrierOneLoop = CUQU_INIT_BARRIER_ONELOOP;
    cuquSafeCall(cudaMalloc(&barrierOneLoop, sizeof (cuqu::detail::barrier_oneloop)));
    cuquSafeCall(cudaMemcpy(barrierOneLoop, &h_barrierOneLoop, sizeof (cuqu::detail::barrier_oneloop), cudaMemcpyHostToDevice));

    cuqu::detail::barrier_lockfree h_barrierLockFree = CUQU_INIT_BARRIER_LOCKFREE;
    cuquSafeCall(cudaMalloc(&barrierLockFree, sizeof (cuqu::detail::barrier_lockfree)));
    cuquSafeCall(cudaMemcpy(barrierLockFree, &h_barrierLockFree, sizeof (cuqu::detail::barrier_lockfree), cudaMemcpyHostToDevice));

    cuqu::detail::barrier_tree_based h_barrierTreeBased = CUQU_INIT_BARRIER_TREE_BASED;
    cuquSafeCall(cudaMalloc(&barrierTreeBased, sizeof (cuqu::detail::barrier_tree_based)));
    cuquSafeCall(cudaMemcpy(barrierTreeBased, &h_barrierTreeBased, sizeof (cuqu::detail::barrier_tree_based), cudaMemcpyHostToDevice));   

    printf("# Kernel calls: %d\n", params->kernelCalls);
    printf("# Test repeats: %d\n", params->testRepeats);
    fflush(stdout);

    // Warm up!
    int maxWorkers = params->numBlocks * params->numThreads;
    printf("# Warming up (blocks=%d, threads=%d, workers=%d)... ", params->numBlocks, params->numThreads, maxWorkers);
    fflush(stdout);
    barrierTime(maxBlocks, maxThreads, maxWorkers, params->kernelCalls, params->testRepeats, params->syncMethod, &measures);
    printf("OK.\n");
    fflush(stdout);

    printf("# Run test.\n");
    printf("#BLK  THR   WRK    AVG_TIME   MIN_TIME   MAX_TIME   ERR_TIME\n"); 
    
    if (params->autoBlocks)  startBlocks  = 1;
    if (params->autoThreads) startThreads = 1;

    for (int currBlocks = startBlocks; currBlocks <= maxBlocks; currBlocks += params->blockStep) {
        for (int currThreads = startThreads; currThreads <= maxThreads; currThreads += params->threadStep) {
            if (params->autoWorkers) {
                for (int currWorkers = currThreads * currBlocks; currWorkers >= 0; currWorkers -= params->workerStep) {
                    printf("%-5d %-5d %-6d ", currBlocks, currThreads, currWorkers);
                    barrierTime(currBlocks, currThreads, currWorkers, params->kernelCalls, params->testRepeats, params->syncMethod, &measures);
                    printf("%-10f %-10f %-10f %-10f\n", measures.averageTime, measures.minTime, measures.maxTime, measures.errorTime);
                    fflush(stdout);
                }
            } else {
                int currWorkers = currBlocks * currThreads;
                printf("%-5d %-5d %-6d ", currBlocks, currThreads, currWorkers);
                barrierTime(currBlocks, currThreads, currWorkers, params->kernelCalls, params->testRepeats, params->syncMethod, &measures);
                printf("%-10f %-10f %-10f %-10f\n", measures.averageTime, measures.minTime, measures.maxTime, measures.errorTime);
                fflush(stdout);
            }
        }
    }

    cuquSafeCall(cudaFree(barrierOneLoop));
    cuquSafeCall(cudaFree(barrierLockFree));
    cuquSafeCall(cudaFree(barrierTreeBased));
}


// Main function
int main(int argc, char *argv[]) {
    
    UserParameters params;

    // Set program name
    programName = basename(argv[0]);
    
    // Set default parameters
    params.autoBlocks  = false;
    params.autoThreads = false;
    params.autoWorkers = false;
    params.numBlocks   = DEFAULT_NUM_BLOCKS;
    params.numThreads  = DEFAULT_NUM_THREADS;
    params.numWorkers  = DEFAULT_NUM_WORKERS;
    params.blockStep   = DEFAULT_BLOCK_STEP;
    params.threadStep  = DEFAULT_THREAD_STEP;
    params.workerStep  = DEFAULT_WORKER_STEP;
    params.testRepeats = DEFAULT_TEST_REPEATS;
    params.kernelCalls = DEFAULT_KERN_CALLS;
    params.syncMethod  = DEFAULT_SYNC_METHOD;

    // Parse arguments and run test
    parseArguments(argc, argv, &params);  
    runTest(&params);        
   
    exit(EXIT_SUCCESS);
}

////////////////////////////////////////////////////////////////////////////////
//
// Utilities.
//
////////////////////////////////////////////////////////////////////////////////

// Compute average and standard deviation.
void getMeasures(float times[], size_t n, Measures *m) {

    float sum = 0, sum2 = 0;

    m->minTime = FLT_MAX;
    m->maxTime = FLT_MIN;

	for (int i = 0; i < n; i++) {
		sum += times[i];
        // Preserve min and max values
        if (m->minTime > times[i]) m->minTime = times[i];
        if (m->maxTime < times[i]) m->maxTime = times[i];
	}

    // Compute average
    m->averageTime = sum / n;

	for (int i = 0; i < n; i++) {
        float v = times[i] - m->averageTime;
		sum2 += v * v;
	}

    // Compute standard deviation
	m->errorTime = sqrt(sum2 / (n - 1));
}


// Print program usage
void printUsage(bool hint) {

    fprintf(stderr, "Usage: %s [OPTIONS]...\n", programName);

    if (hint) {
        fprintf(stderr, "Try `%s --help' for more information.\n", programName);
    }
    else {
        fprintf(stderr,
            "\nOPTIONS:\n"\
            "      --auto-blocks         auto increment blocks\n" \
            "      --auto-threads        auto increment threads\n" \
            "      --auto-workers        auto increment workers\n" \
            "  -b, --blocks=NUM          use NUM blocks (default=%d)\n" \
            "  -B, --block-step=NUM      incremental step of blocks\n" \
            "  -h, --help                display this help and exit\n" \
            "  -k, --kernel-calls=NUM    call kernel NUM times (default=%d)\n" \
            "  -r, --test-repeats=NUM    repeat test NUM times (default=%d)\n" \
            "  -s, --sync-method=METHOD  use METHOD barrier. Possible values:\n" \
            "                            nosync, cpu_implicit, cpu_explicit, \n" \
	        "                            oneloop, lockfree, tree_based\n" \
            "  -t, --threads=NUM         use NUM threads (default=%d)\n" \
            "  -T, --thread-step=NUM     incremental step of threads\n" \
            "  -w, --workers=NUM         use NUM workers (default=%d)\n" \
            "  -W, --worker-step=NUM     incremental step of working threads\n" \

            "\nReport bugs to <marco.esposito@gmail.com>.\n",
            DEFAULT_NUM_BLOCKS, 
            DEFAULT_BLOCK_STEP,
            DEFAULT_KERN_CALLS,
            DEFAULT_TEST_REPEATS,
            DEFAULT_NUM_THREADS,
            DEFAULT_THREAD_STEP,
            DEFAULT_NUM_WORKERS,
            DEFAULT_WORKER_STEP
        );
    }
}

// Parse an argument number. Number must be greater than 0
void parseNumberGTZ(char opt, char *optarg, int *numptr) {
    
    char *endptr;

    // Convert string to number
    *numptr = (int) strtol(optarg, &endptr, 10);
                
    // Must be a valid number greater than 0
    if (optarg[0] == '\0' || endptr[0] != '\0' || *numptr <= 0) {   
        fprintf(stderr, "The argument of option `%c' must be a number greater than 0.\n", opt);  
        printUsage(true);
        exit(EXIT_FAILURE);
    }
}

// Parse program arguments
void parseArguments(int argc, char *argv[], UserParameters *params) {

    while (1) {
        static struct option longOptions[] = 
        {
            {"auto-blocks",  no_argument,       NULL, 0},
            {"auto-threads", no_argument,       NULL, 0},            
            {"auto-workers", no_argument,       NULL, 0},            

            {"block-step",   required_argument, NULL, 'B'},
            {"sync-method",  required_argument, NULL, 's'}, 
            {"thread-step",  required_argument, NULL, 'T'},            
            {"worker-step",  required_argument, NULL, 'W'},        

            {"blocks",       required_argument, NULL, 'b'},            
            {"threads",      required_argument, NULL, 't'},            
            {"workers",      required_argument, NULL, 'w'},            

            {"help",         no_argument,       NULL, 'h'},

            {"kernel-calls", required_argument, NULL, 'k'},
            {"test-repeats", required_argument, NULL, 'r'},
            {NULL,           no_argument,       NULL,  0 }
        };

        // getopt_long stores the option index here 
        int optionIndex = 0;        
        int c = getopt_long(argc, argv, "B:T:W:b:t:w:hk:r:s:", longOptions, &optionIndex);

        // Detect the end of the options
        if (c == -1) {
            break;
        }

        switch (c) {

            case 'h':
                printUsage(0);
                exit(EXIT_SUCCESS);
                break;

            // Auto flags
            case 0:
                // If this option set a flag, do nothing else now
                if (longOptions[optionIndex].flag != 0) {
                    break;
                }
                
                if (strcmp("auto-blocks", longOptions[optionIndex].name) == 0) {
                    params->autoBlocks = true;
                } else if (strcmp("auto-threads", longOptions[optionIndex].name) == 0) {
                    params->autoThreads = true;
                } else if (strcmp("auto-workers", longOptions[optionIndex].name) == 0) {
                    params->autoWorkers = true;  
                } 
            break;

            // Main params
            case 'b':
                parseNumberGTZ(c, optarg, &(params->numBlocks));
                break;

            case 't':
                parseNumberGTZ(c, optarg, &(params->numThreads));
                break;
           
            case 'w':
                parseNumberGTZ(c, optarg, &(params->numWorkers));
                break;            

            // Steps
            case 'B':
                parseNumberGTZ(c, optarg, &(params->blockStep));
                break;

            case 'T':
                parseNumberGTZ(c, optarg, &(params->threadStep));
                break;
           
            case 'W':
                parseNumberGTZ(c, optarg, &(params->workerStep));
                break;

            case 's':
                // Select sync implementation method
                if (strcmp("nosync", optarg) == 0) {
                    params->syncMethod = NoSync;
                } else if (strcmp("oneloop", optarg) == 0) {
                    params->syncMethod = OneLoop;
                } else if (strcmp("lockfree", optarg) == 0) {
                    params->syncMethod = LockFree;
                } else if (strcmp("tree_based", optarg) == 0) {
                    params->syncMethod = TreeBased;
                } else if (strcmp("cpu_implicit", optarg) == 0) {
                    params->syncMethod = CPUImplicit;
                } else if (strcmp("cpu_explicit", optarg) == 0) {
                    params->syncMethod = CPUExplicit;
                }

                break;

            // Loops
            case 'k':
                parseNumberGTZ(c, optarg, &(params->kernelCalls));
                break;

            case 'r':
                parseNumberGTZ(c, optarg, &(params->testRepeats));
                break;

            case '?':
                // getopt_long already printed an error message
                printUsage(true);
                exit(EXIT_FAILURE);

            default:
                abort();
        }
    }
}
